{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install langchain milvus pymilvus python-dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.memory import VectorStoreRetrieverMemory\n",
    "from langchain.chains import ConversationChain\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "import openai\n",
    "load_dotenv()\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from milvus import default_server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "    __  _________ _   ____  ______\n",
      "   /  |/  /  _/ /| | / / / / / __/\n",
      "  / /|_/ // // /_| |/ / /_/ /\\ \\\n",
      " /_/  /_/___/____/___/\\____/___/ {Lite}\n",
      "\n",
      " Welcome to use Milvus!\n",
      "\n",
      " Version:   v2.2.8-lite\n",
      " Process:   13978\n",
      " Started:   2023-05-30 08:34:48\n",
      " Config:    /Users/yujiantang/.milvus.io/milvus-server/2.2.8/configs/milvus.yaml\n",
      " Logs:      /Users/yujiantang/.milvus.io/milvus-server/2.2.8/logs\n",
      "\n",
      " Ctrl+C to exit ...\n"
     ]
    }
   ],
   "source": [
    "default_server.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Milvus\n",
    "embeddings = OpenAIEmbeddings()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install pymilvus\n",
    "from pymilvus import utility, connections\n",
    "\n",
    "connections.connect(host=\"127.0.0.1\", port=default_server.listen_port)\n",
    "utility.drop_collection('LangChainCollection')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "vectordb = Milvus.from_documents(\n",
    "    {},\n",
    "    embeddings,\n",
    "    connection_args={\"host\": \"127.0.0.1\", \"port\": default_server.listen_port})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = Milvus.as_retriever(vectordb, search_kwargs=dict(k=1))\n",
    "memory = VectorStoreRetrieverMemory(retriever=retriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "about_me = [\n",
    "    {\"input\": \"My favorite snack is chocolate\",\n",
    "     \"output\": \"Nice\"},\n",
    "    {\"input\": \"My favorite sport is swimming\",\n",
    "     \"output\": \"Cool\"},\n",
    "    {\"input\": \"My favorite beer is Guinness\",\n",
    "     \"output\": \"Great\"},\n",
    "    {\"input\": \"My favorite dessert is cheesecake\",\n",
    "     \"output\": \"Good to know\"},\n",
    "    {\"input\": \"My favorite musician is Taylor Swift\",\n",
    "     \"output\": \"Same\"}\n",
    "]\n",
    "for example in about_me:\n",
    "    memory.save_context({\"input\": example[\"input\"]}, {\"output\": example[\"output\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: My favorite musician is Taylor Swift\n",
      "output: Same\n"
     ]
    }
   ],
   "source": [
    "print(memory.load_memory_variables({\"prompt\": \"who is my favorite musician?\"})[\"history\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: My favorite snack is chocolate\n",
      "output: Nice\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi, my name is Gary, what's up?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi Gary, my name is AI. I'm doing well. How about you?\""
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm = OpenAI(temperature=0) # Can be any valid LLM\n",
    "_DEFAULT_TEMPLATE = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
    "\n",
    "Relevant pieces of previous conversation:\n",
    "{history}\n",
    "\n",
    "(You do not need to use these pieces of information if not relevant)\n",
    "\n",
    "Current conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "PROMPT = PromptTemplate(\n",
    "    input_variables=[\"history\", \"input\"], template=_DEFAULT_TEMPLATE\n",
    ")\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm, \n",
    "    prompt=PROMPT,\n",
    "    memory=memory,\n",
    "    verbose=True\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, my name is Gary, what's up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: who is my favorite artist?\n",
      "response:  You said your favorite musician is Taylor Swift.\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: who is my favorite musician?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' You said your favorite musician is Taylor Swift. Is there anything else I can help you with?'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"who is my favorite musician?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Retrying langchain.llms.openai.completion_with_retry.<locals>._completion_with_retry in 4.0 seconds as it raised RateLimitError: The server had an error while processing your request. Sorry about that!.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: My favorite dessert is cheesecake\n",
      "output: Good to know\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: Whats my favorite dessert?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' You said your favorite dessert is cheesecake.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"Whats my favorite dessert?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: Hi, my name is Gary, what's up?\n",
      "response:  Hi Gary, my name is AI. I'm doing well. How about you?\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: What's my name?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' You said your name is Gary.'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"What's my name?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "default_server.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hw_milvus",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
